package com.yupi.yuaiagent.rag;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.markdown.MarkdownDocumentReader;
import org.springframework.ai.reader.markdown.config.MarkdownDocumentReaderConfig;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.ResourcePatternResolver;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

@Component
@Slf4j
public class LoveAppDocumentLoader {

    private final ResourcePatternResolver resourcePatternResolver;;

    public LoveAppDocumentLoader(ResourcePatternResolver resourcePatternResolver) {
        this.resourcePatternResolver = resourcePatternResolver;
    }

    public List<Document> loadMarkdowns() throws IOException {
        List<Document> documents = new ArrayList<>();
        // 获取所有md文件
        Resource[] resources = resourcePatternResolver.getResources("classpath:document/*.md");
        for (Resource resource : resources) {
            String fileName = resource.getFilename();
            // 提取文档倒数第 3 和第 2 个字作为标签
            String status = fileName.substring(fileName.length() - 6, fileName.length() - 4);
            MarkdownDocumentReaderConfig config = MarkdownDocumentReaderConfig.builder()
                    .withHorizontalRuleCreateDocument(true)  // 当设置为 true 时，Markdown中的水平规则将创建新的 Document 对象。
                    .withIncludeCodeBlock(false)  // 当设置为 true 时，代码块将包含在与周围文本相同的 Document 中。当 false 时，代码块创建单独的 Document 对象。
                    .withIncludeBlockquote(false)  // 当设置为 true 时，blockquote将包含在与周围文本相同的 Document 中。当 false 时，blockquote创建单独的 Document 对象。
                    .withAdditionalMetadata("filename", fileName)  // 允许您添加自定义元数据到所有创建的 Document 对象。
                    .withAdditionalMetadata("status", status)
                    .build();
            MarkdownDocumentReader reader = new MarkdownDocumentReader(resource, config);
            documents.addAll(reader.get());
        }
        return documents;
    }

}
